package assignment01;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Set;


import nlp.langmodel.LanguageModel;
import nlp.util.Counter;
import nlp.util.Pair;

public class InterpolatedBigramLanguageModel implements LanguageModel {
	private static final int GENERATED_SENTENCE_LENGTH = 25;
	private static String START = "<S>";
	private static String STOP = "</S>";
	private static double INCREMENT_FACTOR = 1.0;
	private GoodTuringCounter<String> unigramCounter = new GoodTuringCounter<String>();
//	private Counter<String> unigramCounter = new Counter<String>();
	private Counter<String> specialCounter = new Counter<String>();
	private GoodTuringCounter<Bigram<String, String>> bigramCounter = new GoodTuringCounter<Bigram<String,String>>();
	private Collection<List<String>> validationCollection;
	Pair<Double, Double> lambdaPair = null;
	private double unigramTotal;
	private double bigramTotal;
	
	public InterpolatedBigramLanguageModel(Collection<List<String>> trainingCollection, Collection<List<String>> validationCollection) {
		this.validationCollection = validationCollection;
		
		
		for (List<String> sentence: trainingCollection) {
			List<String> stoppedSentence = new ArrayList<String>(sentence);
			stoppedSentence.add(0, START);
			stoppedSentence.add(STOP);
			
			// unigram generation
			for (String word : stoppedSentence) {
				unigramCounter.incrementCount(word, INCREMENT_FACTOR);
			}
			
			// bigram generation
			Iterator<String> iter = stoppedSentence.iterator();
			String first = iter.next();
			String second = iter.next();
			while (iter.hasNext()) {
				Bigram<String, String> bigram = new Bigram<String, String>(first, second);
				bigramCounter.incrementCount(bigram, INCREMENT_FACTOR);
				first = second;
				if (iter.hasNext()) {
					second = iter.next();
				}
			}
		}
		unigramCounter.smooth();

		bigramCounter.smooth();
		
		// special counter
		for (Bigram<String, String> bigram : bigramCounter.keySet()) {
				specialCounter.incrementCount(bigram.first, INCREMENT_FACTOR * bigramCounter.getCount(bigram));
		}
		
		// nomarlizing counter
//		unigramCounter.smooth();
		
		unigramTotal = unigramCounter.totalCount();
		bigramTotal = bigramCounter.totalCount();
		System.out.println("uniCount = " + unigramTotal);
		System.out.println("bigramCount = " + bigramTotal);
		System.out.println("unigram vacabulary size is " + unigramCounter.keySet().size());
		System.out.println("bigram vacabulary size is " + bigramCounter.keySet().size());
		System.out.println("special size is " + specialCounter.size());
		
		// train validation set
//		lambdaPair = resolveLambda();
		lambdaPair = new Pair<Double, Double>(0.5, .5);
		System.out.println("lambdaPair is " + lambdaPair);
		System.out.println("LanguageModelTester.calculatePerplexity() is " + LanguageModelTester.calculatePerplexity(this, validationCollection));
	}
	

	
	private Pair<Double, Double> resolveLambda() {
		double gap = .01;
		int times = (int)(1.0 / gap);
		
		double lambda1 = 0.0;
		double lambda2 = 1.0 - lambda1;
		double perplexity = Double.POSITIVE_INFINITY;
		for (int i = 1; i < times; i++) {
			double temp = calculatePerplexity(validationCollection, i * gap, 1.0 - i * gap);
			if (temp < perplexity) {
				perplexity = temp;
				lambda1 = gap * i;
				lambda2 = 1.0 - lambda1;
			}
//			System.out.println("temp perplexity of rount " + i + " is " + temp);
		}
		System.out.println("\nlowest perplexity " + perplexity);
		return new Pair<Double, Double>(lambda1, lambda2);
	}
	
	private double getInterpolatedWordLogProbability(List<String> sentence, int index, double lambda1, double lambda2) {
		if (index <= 0) {
			return 0.0;
		}
		
		String current = sentence.get(index);
		String previous = sentence.get(index - 1);
		
		Bigram<String, String> key = new Bigram<String, String>(previous, current);
		
		double uniCount = specialCounter.getCount(previous);
//		for (Bigram<String, String> bigram : bigramCounter.keySet()) {
//			if (bigram.first.equals(previous)) {
//				uniCount += bigramCounter.getCount(bigram);
//			}
//		}
		
		double pUnigram = 0.0;
		if (!unigramCounter.keySet().contains(current)) {
//			pUnigram = unigramCounter.getZeroCount() / unigramTotal;
		} else {
			pUnigram = unigramCounter.getCount(current) / unigramTotal;
		}
		
		double pBigram = 0.0;
		if (!bigramCounter.keySet().contains(key)) {
			pBigram = bigramCounter.getZeroCount() / bigramTotal;
		} else {
			pBigram = bigramCounter.getCount(key) / uniCount;
		}
		
		double logProbability1 = Double.isInfinite(Math.log(lambda1)) ? Double.NEGATIVE_INFINITY : Math.log(lambda1) + Math.log(pUnigram);
		double logProbability2 = Double.isInfinite(Math.log(lambda2)) ? Double.NEGATIVE_INFINITY : Math.log(lambda2) + Math.log(pBigram);
		
//		System.out.println("p1 is " + logProbability1 + ", and p2 is " + logProbability2);
		double temp = Math.pow(Math.E, logProbability1) + Math.pow(Math.E, logProbability2);
		return  Math.log(temp) / Math.log(2.0);
	}
	

	public double getInterplatedSentenceLogProbability(List<String> sentence, double lambda1, double lambda2) {
		double logProbability = 0.0;
		List<String> stoppedSentence = new ArrayList<String>(sentence);
		stoppedSentence.add(0, START);
		stoppedSentence.add(STOP);
		for (int i = 0; i < stoppedSentence.size(); i++) {
			logProbability += getInterpolatedWordLogProbability(stoppedSentence, i, lambda1, lambda2);
		}
//		System.out.println("sentence p is " + logProbability);
		return logProbability;
	}
	
	private double calculatePerplexity(Collection<List<String>> sentenceCollection, double lambda1, double lambda2) {
	    double logProbability = 0.0;
	    double numSymbols = 0.0;
	    for (List<String> sentence : sentenceCollection) {
	      logProbability += getInterplatedSentenceLogProbability(sentence, lambda1, lambda2);
	      numSymbols += sentence.size();
	    }
	    double avgLogProbability = logProbability / numSymbols;
	    double perplexity = Math.pow(.5, avgLogProbability);
	    return perplexity;
	  }
	
	@Override
	public double getSentenceProbability(List<String> sentence) {
		double probability = 0.0;
		List<String> stoppedSentence = new ArrayList<String>(sentence);
		stoppedSentence.add(0, START);
		stoppedSentence.add(STOP);
		
		for (int index = 0; index < stoppedSentence.size(); index++) {
			probability += getInterpolatedWordLogProbability(stoppedSentence, index, lambdaPair.getFirst().doubleValue(), lambdaPair.getSecond().doubleValue());
		}
//		System.out.println("sentence probability is " + Math.pow(2.0, probability));
		return Math.pow(2.0, probability);
	}
/**
	private String generateWord() {
		 double sample = Math.random();
		 double sum = 0.0;
//		 double sum = bigramCounter.getZeroCount() / bigramCounter.totalCount();
		 for (Bigram<String, String> key : bigramCounter.keySet()) {
			 
			 double p = bigramCounter.getCount(key) / unigramCounter.getCount(key.first);
			 sum += p;
			 System.out.println("sum is " + sum + ", and sample is " + sample);
			 if (sum > sample) {
				 return key.second;
			 }
		 }
//		return "*UNKNOWN*";
		 return STOP;
	}
	*/
	private String generateWord(String previous) {
		double sample = Math.random();
		double sum = 0.0;
		double uniCount = specialCounter.getCount(previous); //unigramCounter.getCount(previous);

//		double sum1 = 0.0;
//		for (Bigram<String, String> bigram : bigramCounter.keySet()) {
//			if (bigram.first.equals(previous)) {
//				sum1 += bigramCounter.getCount(bigram) ;
//			}
//		}
//		System.out.println("sum should actually be " + sum1 + ", while uniCount is " + uniCount);
		
		for (Bigram<String, String> bigram : bigramCounter.keySet()) {
			if (bigram.first.equals(previous)) {
				sum += bigramCounter.getCount(bigram) / uniCount;
			}
			if (sum > sample) {
				return bigram.second;
			}
		}		
		
		return STOP;
	}
	
	@Override
	public List<String> generateSentence() {
		List<String> sentence = new ArrayList<String>();
		sentence.add(START);
		String word = generateWord(START);
		while (!word.equals(STOP) && sentence.size() < GENERATED_SENTENCE_LENGTH) {
//			System.out.println(word);
			sentence.add(word);
			word = generateWord(word);
		}
		sentence.add(STOP);
		System.out.println("sentence size is " + sentence.size());
		return sentence;
	}
}
